# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.

import torch

from atk.configs.dataset_config import InputDataset
from atk.tasks.api_execute import register
from atk.tasks.api_execute.base_api import BaseApi
from atk.tasks.dataset.base_dataset import OpsDataset
from atk.configs.results_config import TaskResult


@register("ascend_aclnn_moe_token_unpermute_probs_none")
class FunctionUnpermuteApi(BaseApi):
    def __init__(self, task_result: TaskResult):
        super(FunctionUnpermuteApi, self).__init__(task_result)
        self.permuted_tokens = None
        self.origin_sorted_indices = None
        self.probs = None

    def __call__(self, input_data: InputDataset, with_output: bool = False):
        orgin_dtype = self.permuted_tokens.dtype
        permuted_tokens = self.permuted_tokens.to(torch.float).cpu()
        sorted_indices = self.origin_sorted_indices.cpu()

        if self.probs is not None:
            probs = self.probs.cpu()
            num_unpermuted_tokens = probs.numel()
            topk = probs.size(1)
            probs = probs.to(torch.float)
        else:
            probs = None
            num_unpermuted_tokens = permuted_tokens.size(0)
            topk = 1

        unpermuted_tokens = torch.zeros(
            [num_unpermuted_tokens, permuted_tokens.shape[-1]],
            dtype=torch.float,
            device=permuted_tokens.device,
        )

        unpermuted_tokens.index_copy_(0, sorted_indices, permuted_tokens)
        unpermuted_tokens = unpermuted_tokens.reshape(-1, topk, permuted_tokens.size(-1))

        # if probs is not None:
        #     unpermuted_tokens = unpermuted_tokens * probs.unsqueeze(-1)
        unpermuted_tokens = unpermuted_tokens.sum(dim=1)

        return unpermuted_tokens.to(orgin_dtype)

    def init_by_input_data(self, input_data: InputDataset):
        tokens = input_data.kwargs["permutedTokens"]
        print(tokens[0])
        indices = input_data.kwargs['sorted_indices']
        print(indices[0])
        if 'probs' in input_data.kwargs:
            probs = input_data.kwargs["probs"]
            topk = probs.shape[1]
        else:
            probs = None
            topk = 1
        totalLen = tokens.shape[0]
        # 生成indices
        indices = torch.randint(low=0, high=topk, size=(int(totalLen/topk), topk)).to(torch.int32)
        # 生成unpermute输入
        permuted_tokens, sorted_indices_tmp = self.permute(tokens, indices)
        input_data.kwargs["permutedTokens"] = permuted_tokens
        input_data.kwargs['sorted_indices'] = torch.argsort(sorted_indices_tmp, stable=True).to(torch.int32)

        print(input_data.kwargs["permutedTokens"][0])
        print(input_data.kwargs['sorted_indices'][0])
        self.permuted_tokens = permuted_tokens
        self.origin_sorted_indices = sorted_indices_tmp.to(torch.int64)
        self.probs = probs
        OpsDataset.seed_everything()

    def permute(self, tokens, indices, num_out_tokens: int = None, padded_mode: bool = False):
        tokens = tokens.cpu()
        indices = indices.cpu()
        if indices.dim() == 1:
            topk = 1
        else:
            topk = indices.size(1)

        flatten_indices = indices.view(-1)
        sorted_indices = torch.argsort(flatten_indices, stable=True)
        tmp = sorted_indices
        if num_out_tokens is not None:
            sorted_indices = sorted_indices[:num_out_tokens]
        permuted_tokens = tokens.index_select(0, sorted_indices // topk)
        return permuted_tokens, tmp


